
from langchain.chat_models import ChatOpenAI
from langchain.chains.summarize import load_summarize_chain
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.prompts import PromptTemplate
from langchain.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    AIMessagePromptTemplate,
    HumanMessagePromptTemplate,
)
from langchain.schema import (
    AIMessage,
    HumanMessage,
    SystemMessage
)


MODEL = "gpt-3.5-turbo"
API_KEY = ""


with open('prompt/prompt.txt') as f:
    prompt_cmd = f.read()

with open('prompt/summarize.txt') as f:
    summarize_cmd = f.read()

with open('prompt/refine_summarize.txt') as f:
    refine_summarize_cmd = f.read()



with open('test_data/news1.txt') as f:
    tar_content = f.read()

print(len(tar_content))

#文档分割器
text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(chunk_size=500, chunk_overlap=0)
tar_docs = text_splitter.create_documents([tar_content])
print(len(tar_docs))
print(tar_docs)

#载入模板
PROMPT = PromptTemplate(template=summarize_cmd, input_variables=["text"])
refine_prompt = PromptTemplate(template=refine_summarize_cmd, input_variables=["existing_answer", "text"])

#定义模型
llm = ChatOpenAI(model_name=MODEL, openai_api_key=API_KEY, temperature=0)

#组件总结链
sum_chain = load_summarize_chain(llm, chain_type="refine", question_prompt=PROMPT, refine_prompt=refine_prompt)

#设置系统信息
sys_msg = SystemMessage(content='You are a stock analyst.')



if len(tar_docs) > 1:
    #总结长文本
    sum_res = sum_chain.run(tar_docs)
    print(sum_res)
    #人类提问
    human_msg = HumanMessage(content=prompt_cmd % sum_res)
    #大模型交互
    chat_res = llm([sys_msg, human_msg])    
    print(chat_res)
else:
    human_msg = HumanMessage(content=prompt_cmd % tar_content)
    chat_res = llm([sys_msg, human_msg])    
    print(chat_res)
    



